Summary

This PR resolves critical execution failures, API deprecations, memory allocation errors, and variable mismatches in the DPO training notebook. The changes ensure reliable, error-free end-to-end training while maximizing accelerator resource efficiency.

Key Changes

1. Runtime & Logic Fixes

• Mesh Sharding: Set mesh axis types to jax.sharding.AxisType.Auto to resolve NamedSharding ValueError.
• Missing Attributes: Applied fixes from PR Fix AttributeError with the SentencePieceProcessor during tokenizer initialization #1719 to prevent attribute errors during execution.
• Dataset Epoch Handling: Added .repeat(EPOCHS) to train_ds to prevent early training termination after epoch 1.
• Imports & Variable Naming: Fixed NameError by importing orbax.checkpoint as ocp, and aligned trainer calls to dpo_trainer.train().
• Image Processing: Removed redundant image_processor from DPOTrainer initialization to fix image array conversion TypeError.
• Logging Dependencies: Fixed TensorBoard logging conflicts by aligning package dependencies and logging API calls.

2. Resource & Performance Optimization

• OOM Reduction: Lowered BATCH_SIZE from 2 to 1 to fit within memory constraints and eliminate JAX RESOURCE_EXHAUSTED allocation failures in jit_compute_logps.

3. API Deprecation Upgrades

• Mesh Context: Updated with mesh: syntax to with jax.set_mesh(mesh):.
• Orbax Checkpointing: Updated checkpoint options to use modern checkpointing_options with save_decision_policy.

4. Usability & Documentation

• Header Badges: Added interactive Google Colab and Kaggle launch badges to the top of the notebook.